import cv2
import time
import torch
import numpy as np
from pathlib import Path
from ultralytics import YOLO
from collections import deque
from utils.Logger import Logger
from typing import Optional, List, Tuple
from config.message_constants import YOLOMessages

class YOLOController:
    def __init__(self, 
                 model_path: str, 
                 status_callback=None,
                 device: str = 'auto', 
                 conf_threshold: float = 0.45,
                 iou_threshold: float = 0.7,
                 anomaly_classes: List[str] = None,
                 logger: Optional[Logger] = None,
                 module: str = "YOLO"
        ):
        """
        初始化YOLO异常检测器

        参数说明：
        :param model_path: 模型文件路径（必须存在）
        :param status_callback: 状态变化回调函数
        :param device: 计算设备（auto/cpu/cuda:0）
        :param conf_threshold: 检测置信度阈值（0-1）
        :param iou_threshold: 非极大抑制阈值（0-1）
        :param anomaly_classes: 需检测的异常类别列表
        :param logger: 日志记录器实例（需预先初始化）
        :param module: 日志模块名称（用于分类日志）
        """
        # 日志记录器初始化（多个控制器采用同一个日志记录器）
        self.logger = logger
        # 设置该模块名称
        self.modelName = module

        # 模型参数初始化
        self.model = None  # YOLO模型实例
        self._model_loaded = False  # 模型加载状态标志
        self.device = self._select_device(device)   # 自动选择设备
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.anomaly_classes = anomaly_classes
        self.is_running = False # 模型运行状态标志
        self.status_callback = status_callback  # 状态回调函数
        # 模型实例加载
        self._load_model(model_path)
        
        # 状态缓存
        self.last_anomaly_time = None   # 最近异常时间戳
        self.detection_history = deque(maxlen=30)  # 记录最近30帧结果


    def _select_device(self, device: str) -> str:
        """
        自动选择最佳计算设备
        
        逻辑说明：
        1. 当设备设为'auto'时，优先检测CUDA可用性
        2. 若CUDA不可用，自动回退到CPU
        3. 记录设备信息到日志系统
        """
        if device == 'auto':
            if torch.cuda.is_available():
                device = f'cuda:{torch.cuda.current_device()}'
                self.logger.info(YOLOMessages.CUDA_DETECTED.format(torch.cuda.get_device_name()),self.modelName)
            else:
                device = 'cpu'
                self.logger.warning(YOLOMessages.CUDA_NOT_DETECTED,self.modelName)
        return device

    def _load_model(self, model_path: str):
        """
        加载 YOLOv11 模型
        
        步骤说明：
        1. 验证模型文件是否存在
        2. 初始化YOLO模型并配置参数
        3. 验证模型类别与用户定义的匹配性
        4. 处理可能出现的加载错误

        # TODO: 潜在问题：
        - 如果模型文件损坏或格式不兼容，将抛出RuntimeError
        """
        try:
            model_file = Path(model_path)
            if not model_file.exists():
                raise FileNotFoundError(YOLOMessages.MODEL_FILE_NOT_FOUND.format(str(model_path)), self.modelName)
                
            self.logger.info(YOLOMessages.MODEL_LOADING.format(str(model_file.name)), self.modelName)
            self.model = YOLO(model_file)
            
            # 配置模型参数
            self.model.conf = self.conf_threshold
            self.model.iou = self.iou_threshold
            self.model.to(self.device)  # 将模型移动到目标设备
            
            # 验证类别匹配
            model_classes = self.model.names.values()
            missing_classes = [c for c in self.anomaly_classes if c not in model_classes]
            if missing_classes:
                self.logger.error(YOLOMessages.MODEL_CLASS_MISMATCH.format(strmissing_classes),self.modelName)
                raise ValueError(YOLOMessages.MODEL_CLASS_NOT_MATCH)
                
            self.logger.info(YOLOMessages.MODEL_LOADED.format(str(self.model.model.args['imgsz'])), self.modelName)
            self._model_loaded = True

        except Exception as e:
            self.logger.error(YOLOMessages.MODEL_INIT_FAIL.format(str(e)),self.modelName)
            raise RuntimeError(YOLOMessages.MODEL_LOADED_Fail) from e

    def detect(self, frame: np.ndarray) -> Tuple[np.ndarray, bool, dict]:
        """
        执行单帧异常检测
        
        处理流程：
        1. 输入检查（运行状态和帧有效性）
        2. 执行模型推理
        3. 解析检测结果并绘制标注
        4. 更新检测历史状态
        5. 生成检测报告

        参数说明：
        :param frame: 输入图像帧（RGB格式）
        :return: (标注后的图像, 是否检测到异常, 检测报告)
        """
        # 输入有效性检查
        if not self.is_running or frame is None:
            return frame, False, {}

        try:
            # 执行模型推理
            results = self.model.predict(
                source=frame,
                imgsz=self.model.model.args['imgsz'],
                verbose=False
            )
            
            # 解析检测结果
            processed_frame, anomalies = self._process_detections(frame, results)
            
            # 更新状态历史
            self._update_detection_history(anomalies)
            
            # 生成检测报告
            report = self._generate_report(results)
            
            return processed_frame, anomalies, report
            
        except Exception as e:
            self.logger.error(YOLOMessages.DETECTION_ERROR.format(str(e)),self.modelName)
            if self.status_callback:
                self.status_callback("DETECTION_ERROR")
            return frame, False, {}

    def _process_detections(self, frame: np.ndarray, results) -> Tuple[np.ndarray, bool]:
        """
        处理模型输出并绘制可视化标注
        
        实现细节：
        - 遍历所有检测框
        - 筛选出预定义的异常类别
        - 在图像上绘制红色边界框和类别标签

        返回说明：
        :return: (标注后的图像, 是否检测到异常)
        """
        has_anomaly = False
        annotated_frame = frame.copy()  # 创建副本避免修改原图
        
        for result in results:
            # 提取检测信息（转换为CPU上的numpy数组）
            boxes = result.boxes.xyxy.cpu().numpy()
            classes = result.boxes.cls.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            
            # 遍历每个检测结果
            for box, cls_id, conf in zip(boxes, classes, confidences):
                cls_name = self.model.names[int(cls_id)]    # 获取类别名称
                
                # 仅处理定义的异常类别
                if cls_name in self.anomaly_classes:
                    has_anomaly = True  # 标记存在异常
                    
                    # 解析边界框坐标（整数化）
                    x1, y1, x2, y2 = map(int, box)

                    # 绘制红色边界框（RGB格式颜色）
                    color = (255, 0, 0)
                    cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
                    
                    # 绘制类别标签
                    label = f"{cls_name} {conf:.2f}"
                    cv2.putText(annotated_frame, label, 
                               (x1, y1 - 10), 
                               cv2.FONT_HERSHEY_SIMPLEX, 
                               0.6, color, 2)
        
        return annotated_frame, has_anomaly

    def _update_detection_history(self, current_anomaly: bool):
        """
        更新检测历史状态并触发持续报警
        
        算法说明：
        - 维护一个最大长度为30的队列记录最近检测结果
        - 当最近5帧中有3帧检测到异常时触发持续报警
        """
        self.detection_history.append(current_anomaly)
        
        # 持续报警机制：最近5帧中有3帧检测到异常
        if sum(self.detection_history) >= 3:
            if not self.last_anomaly_time:
                self.last_anomaly_time = time.time()    # 记录首次触发时间
                self.logger.warning(YOLOMessages.PERSISTENT_ANOMALY,self.modelName)
                if self.status_callback:
                    self.status_callback("ANOMALY_PERSISTENT")
        else:
            self.last_anomaly_time = None   # 重置状态

    def _generate_report(self, results) -> dict:
        """
        生成结构化检测报告
        
        报告内容：
        - 时间戳
        - 异常详细信息（类别、数量、最高置信度）
        - 统计信息（总检测数、类别分布）
        """
        report = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "anomalies": [],
            "statistics": {
                "total_detections": 0,
                "class_distribution": {}
            }
        }
        
        # 遍历所有检测结果
        for result in results:
            # 统计每个类别
            for cls_id in result.boxes.cls.unique():
                cls_name = self.model.names[int(cls_id)]
                count = (result.boxes.cls == cls_id).sum().item()

                # 更新统计信息
                report["statistics"]["class_distribution"][cls_name] = count
                report["statistics"]["total_detections"] += count
                
                # 记录异常详情
                if cls_name in self.anomaly_classes:
                    report["anomalies"].append({
                        "class": cls_name,
                        "count": count,
                        "max_confidence": result.boxes.conf[result.boxes.cls == cls_id].max().item()
                    })
        
        return report

    def enable(self):
        """启用检测器"""
        if not self.is_running:
            self.is_running = True
            self.logger.info(YOLOMessages.DETECTION_ENABLED,self.modelName)
            if self.status_callback:
                self.status_callback("DETECTION_ACTIVE")

    def disable(self):
        """停用检测功能并清空历史状态"""
        if self.is_running:
            self.is_running = False
            self.detection_history.clear()
            self.logger.info(YOLOMessages.DETECTION_DISABLED,self.modelName)
            if self.status_callback:
                self.status_callback("DETECTION_IDLE")

    def __del__(self):
        """析构函数，确保资源正常释放"""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        if self.model is not None:
            try:
                del self.model
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                self.logger.info(YOLOMessages.MODEL_RELEASED,self.modelName)
            except Exception as e:
                self.logger.error(YOLOMessages.RELEASE_FAIL.format(str(e)),self.modelName)